import random
import time

import numpy as np

from src.impl.output import get_output_float16
from src.generator_utils import get_random_seed_tensor
from src.evaluator import has_precision_problem_occurred
from src.classification.classification import get_benchmark


corpus_tf = []
corpus_torch = []
corpus_mnn = []
max_dis_tf = 0
max_dis_torch = 0
max_dis_mnn = 0

sum_cross_frameworks = 0


def dis(tensor, operator, framework, mode='1'):

    benchmark, variable = get_benchmark(tensor, operator)

    tf_output_16, torch_output_16, mnn_output_16, _ = get_output_float16(tensor, operator, variable)

    if framework == 'tf':
        output = tf_output_16.astype('float16')
    elif framework == 'torch':
        output = torch_output_16.astype('float16')
    else:
        output = mnn_output_16.astype('float16')

    if mode == '1':
        diff = np.maximum(benchmark - output, output - benchmark)
        dis = np.sum(diff) / np.size(diff)
        return dis
    elif mode == 'n':
        diff = np.maximum(benchmark - output, output - benchmark)
        return np.max(diff)


def mutation(tensor):

    random_int = random.randint(0,2)
    if random_int == 0:
        theta = 1e-4
    elif random_int == 1:
        theta = 1e-6
    else:
        theta = 1e-8
    disturb = random.uniform(0, theta)

    return tensor + disturb


if __name__ == '__main__':

    shape = tuple(sys.argv[1])
    b = float(sys.argv[2])
    operator = sys.argv[3]
    mode = '1'

    sum_tf = 0
    sum_torch = 0
    sum_mnn = 0
    t = 0

    time_predoo = time.time()
    while time.time() - time_predoo < int(sys.argv[4]) * 3600:
        if t % 100 == 0:
            print('{0} case.'.format(t + 1))
            
        if len(corpus_tf) != 0:
            random_int = random.randint(0, len(corpus_tf)-1)
            random_sample = corpus_tf[random_int]
        else:
            random_sample = get_random_seed_tensor(shape)

        mutated_sample = mutation(random_sample)
        if has_precision_problem_occurred(mutated_sample, operator, 'MRE', b)[0]:
            sum_tf += 1
            np.save('./data/tf/{0}.npy'.format(sum_tf-1), mutated_sample)
            sum_cross_frameworks += 1
            if dis(mutated_sample, operator, 'tf', mode) > max_dis_tf:
                max_dis_tf = dis(mutated_sample, operator, 'tf', mode)
                corpus_tf.append(mutated_sample)

        if len(corpus_torch) != 0:
            random_int = random.randint(0, len(corpus_torch)-1)
            random_sample = corpus_torch[random_int]
        else:
            random_sample = get_random_seed_tensor(shape)

        mutated_sample = mutation(random_sample)

        if has_precision_problem_occurred(mutated_sample, operator, 'MRE', b)[1]:
            sum_torch += 1
            np.save('./data/torch/{0}.npy'.format(sum_torch-1), mutated_sample)
            sum_cross_frameworks += 1
            if dis(mutated_sample, operator, 'torch', mode) > max_dis_torch:
                max_dis_torch = dis(mutated_sample, operator, 'torch', mode)
                corpus_torch.append(mutated_sample)

        if len(corpus_mnn) != 0:
            random_int = random.randint(0, len(corpus_mnn)-1)
            random_sample = corpus_mnn[random_int]
        else:
            random_sample = get_random_seed_tensor(shape)

        mutated_sample = mutation(random_sample)
        if has_precision_problem_occurred(mutated_sample, operator, 'MRE', b)[2]:
            sum_mnn += 1
            np.save('./data/mnn/{0}.npy'.format(sum_mnn-1), mutated_sample)
            sum_cross_frameworks += 1
            if dis(mutated_sample, operator, 'mnn', mode) > max_dis_mnn:
                max_dis_mnn = dis(mutated_sample, operator, 'mnn', mode)
                corpus_mnn.append(mutated_sample)
        t += 1
    time_predoo = time.time() - time_predoo
    print(time_predoo)
    print(t)
    print(len(corpus_tf))
    print(len(corpus_torch))
    print(len(corpus_mnn))